#include "ws_msg_dispatcher.hpp"
#include "nlohmann/json.hpp"
#include "p2pms_log.hpp"

#define MINICORO_IMPL
#include "minicoro/minicoro.h"

namespace p2pms 
{

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
WsMsgDispatcher::WsMsgDispatcher() {}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
WsMsgDispatcher::~WsMsgDispatcher() 
{
	Stop();
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
bool WsMsgDispatcher::Init(int timeout_ms, const EventHandler& event_handler) 
{
	if (timeout_ms == 0) {
		LOG_ERR("invalid timeout_ms:{}", timeout_ms);
		return false;
	}

	m_timeout_ms = timeout_ms;
	m_event_handler = event_handler;

	m_ws_client = std::make_unique<WebSocketClient>();

	if (!m_ws_client->Init(this)) {
		LOG_ERR("Failed to initialize WebSocket client");
		return false;
	}

	LOG_INF("WebSocket client initialized successfully, timeout: {} ms", m_timeout_ms);

	return true;
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
bool WsMsgDispatcher::Start(const std::string& uri)
{
	if (!m_ws_client->Start(uri)) {
		LOG_ERR("Failed to start WebSocket client with URI: {}", uri);
		return false;
	}

	m_running = true;
	m_timeout_thread = std::thread([this] { TimeoutThreadFunc(); });

	return true;
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::Stop() 
{
	m_running = false;
	if (m_ws_client) m_ws_client->Stop();
	if (m_timeout_thread.joinable()) m_timeout_thread.join();
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::Register(const std::string& cmd, const MessageHandler& handler) 
{
	if (cmd.empty()) {
		LOG_ERR("Invalid cmd");
		return;
	}

	if (!handler) {
		LOG_ERR("Invalid handler");
		return;
	}

	std::lock_guard<std::mutex> lock(m_mutex);

	if (m_cmd_handlers.find(cmd) != m_cmd_handlers.end()) {
		LOG_ERR("Command {} already registered", cmd);
	}
	else {
		m_cmd_handlers[cmd] = handler;
		LOG_INF("Registered command: {}", cmd);
	}
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::Unregister(const std::string& cmd)
{
	if (cmd.empty()) {
		LOG_ERR("Invalid cmd: {}", cmd);
		return;
	}

	std::lock_guard<std::mutex> lock(m_mutex);

	auto it = m_cmd_handlers.find(cmd);
	if (it != m_cmd_handlers.end()) {
		m_cmd_handlers.erase(it);
		LOG_INF("Unregistered command: {}", cmd);
	} else {
		LOG_ERR("Command {} not found for unregistration", cmd);
	}
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
bool WsMsgDispatcher::SendMsg(const Message& msg, const std::string& res_cmd) 
{
	if (!m_ws_client) {
		LOG_ERR("WebSocket client not initialized");
		return false;
	}

	// 如果有响应cmd，添加到pending列表中
	if (!res_cmd.empty()) {
		auto pending = std::make_shared<Pending>();
		pending->request = msg;
		pending->res_cmd = res_cmd;
		pending->expire = std::chrono::steady_clock::now() + std::chrono::milliseconds(m_timeout_ms);
		pending->type = PendingType::Async;
		{
			std::lock_guard<std::mutex> lock(m_mutex);

			auto pending_key = PendingKey{res_cmd, msg.rid};
			if (m_pendings.find(pending_key) != m_pendings.end()) {
				LOG_ERR("Pending request {} already exists", pending_key.ToString());
				return false; // 已存在相同的请求
			}
			m_pendings[pending_key] = pending;
			LOG_DBG("Added asynchronous pending request for {}", pending_key.ToString());
		}
	}

	// 发送消息
	return m_ws_client->SendMsg(nlohmann::json(msg).dump());
}

//--------------------------------------------------------------------------------------------------
// TODO: 使用 rid 作为 key 不够严谨，后续可以考虑使用 cmd + rid 作为 key
//--------------------------------------------------------------------------------------------------
MsgResult WsMsgDispatcher::SendMsgSync(const Message& msg, const std::string& res_cmd)
{
	if (!m_ws_client) {
		LOG_ERR("WebSocket client not initialized");
		return { Message{}, CBType::Error };
	}

	auto pending_key = PendingKey{res_cmd, msg.rid};

	auto pending = std::make_shared<Pending>();
	pending->request = msg;
	pending->res_cmd = res_cmd;
	pending->expire = std::chrono::steady_clock::now() + std::chrono::milliseconds(m_timeout_ms);
	pending->type = PendingType::Sync;

	{
		std::lock_guard<std::mutex> lock(m_mutex);

		if (m_pendings.find(pending_key) != m_pendings.end()) {
			LOG_ERR("Pending request {} already exists", pending_key.ToString());
			return { Message{}, CBType::Error };
		}
		m_pendings[pending_key] = pending;
		LOG_DBG("Added synchronous pending request for {}", pending_key.ToString());
	}

	if (!m_ws_client->SendMsg(nlohmann::json(msg).dump())) {
		std::lock_guard<std::mutex> lock(m_mutex);
		m_pendings.erase(pending_key);
		LOG_ERR("Failed to send message: {}", msg.cmd);
		return { Message{}, CBType::Error };
	}

	std::unique_lock<std::mutex> lk(pending->mtx);
	auto timeout = std::chrono::milliseconds(m_timeout_ms);
	if (!pending->cv.wait_for(lk, timeout, [&] { return pending->cbt != CBType::Invalid; })) {
		std::lock_guard<std::mutex> lock(m_mutex);
		m_pendings.erase(pending_key);
		return { Message{}, CBType::Timeout };
	}

	Message response = pending->response;
	CBType cbt = pending->cbt;

	{
		std::lock_guard<std::mutex> lock(m_mutex);
		m_pendings.erase(pending_key);
	}

	return { response, cbt };
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
MsgResult WsMsgDispatcher::SendMsgCoro(const Message& msg, const std::string& res_cmd)
{
	// 1. 必须在协程中调用
	if (!mco_running()) {
		LOG_ERR("SendMsgCoro must be called from within a coroutine");
		return { Message{}, CBType::Error };
	}

	// 2. 必须能获取当前 rtc::Thread，用于后续 resume
	rtc::Thread* current_thread = rtc::Thread::Current();
	if (!current_thread) {
		LOG_ERR("SendMsgCoro must be called from a running rtc::Thread");
		return { Message{}, CBType::Error };
	}

	auto pending_key = PendingKey{ res_cmd, msg.rid };

	// 3. 创建 Pending 请求，保存协程句柄和当前线程
	auto pending = std::make_shared<Pending>();
	pending->request = msg;
	pending->res_cmd = res_cmd;
	pending->expire = std::chrono::steady_clock::now() + std::chrono::milliseconds(m_timeout_ms);
	pending->type = PendingType::Coroutine;
	pending->co = mco_running(); // 获取当前正在运行的协程
	pending->thread = current_thread;

	{
		std::lock_guard<std::mutex> lock(m_mutex);
		if (m_pendings.find(pending_key) != m_pendings.end()) {
			LOG_ERR("Pending request {} already exists", pending_key.ToString());
			return { Message{}, CBType::Error };
		}
		m_pendings[pending_key] = pending;
	}

	// 4. 发送消息
	if (!m_ws_client->SendMsg(nlohmann::json(msg).dump())) {
		LOG_ERR("Send message failed");
		std::lock_guard<std::mutex> lock(m_mutex);
		m_pendings.erase(pending_key);
		return { Message{}, CBType::Error };
	}

	// 5. 关键：yield 协程，让出执行权，等待异步回调 resume
	mco_yield(pending->co);

	// ---- 执行到这里时，说明协程已被 OnWsMessage 或 TimeoutThreadFunc resume ----

	// 6. 获取结果并清理
	MsgResult result = { pending->response, pending->cbt };
	{
		std::lock_guard<std::mutex> lock(m_mutex);
		m_pendings.erase(pending_key);
	}
	
	return result;
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::OnWsMessage(const std::string& msg) 
{
	nlohmann::json jmsg = nlohmann::json::parse(msg, nullptr, false);
	if (!jmsg.is_object()) {
		LOG_ERR("Invalid message format: {}", msg);
		return;
	}
	Message message = jmsg.get<Message>();

	// 响应消息处理
	{
		auto pending_key = PendingKey{jmsg.value("cmd", ""), jmsg.value("rid", "")};

		std::lock_guard<std::mutex> lock(m_mutex);
		auto it = m_pendings.find(pending_key);
		if (it != m_pendings.end()) {
			auto pending = it->second;
			if (pending->type == PendingType::Sync) {
				{
					std::lock_guard<std::mutex> lk(pending->mtx);
					pending->response = message;
					pending->cbt = CBType::Normal;
				}
				LOG_DBG("Synchronous response received {}", pending_key.ToString());
				pending->cv.notify_one();
			}
			else if (pending->type == PendingType::Coroutine) {
				// 为协程设置响应结果，并让其线程 resume 协程
				pending->response = message;
				pending->cbt = CBType::Normal;
				if (pending->co && pending->thread) {
					pending->thread->PostTask([co = pending->co]() {
						mco_resume(co);
					});
				}
				// 注意：这里不移除 pending，由协程 resume 后自己移除
			}
			else if (pending->type == PendingType::Async) {
				auto handler_it = m_cmd_handlers.find(pending->res_cmd);
				if (handler_it != m_cmd_handlers.end()) {
					handler_it->second(message, CBType::Normal);
				}
				LOG_DBG("Asynchronous response handled {}", pending_key.ToString());
				it = m_pendings.erase(it);
			} else {
				LOG_ERR("Unknown pending type for key: {}", pending_key.ToString());
			}
			return; // 响应已处理
		}
	}

	// 普通消息分发
	MessageHandler handler;
	{
		std::lock_guard<std::mutex> lock(m_mutex);
		auto it = m_cmd_handlers.find(message.cmd);
		if (it != m_cmd_handlers.end()) {
			handler = it->second;
		}
	}
	if (handler) {
		handler(message, CBType::Normal);
	}
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::TimeoutThreadFunc() 
{
	while (m_running) {
		std::this_thread::sleep_for(std::chrono::milliseconds(100));
		auto now = std::chrono::steady_clock::now();
		std::vector<PendingKey> timeout_pendings;
		{
			std::lock_guard<std::mutex> lock(m_mutex);
			for (const auto& [key, pending] : m_pendings) {
				if (now > pending->expire) {
					if (pending->type == PendingType::Sync) {
						{
							std::lock_guard<std::mutex> lk(pending->mtx);
							pending->cbt = CBType::Timeout;
						}
						pending->cv.notify_one();
					}
					else if (pending->type == PendingType::Coroutine) {
						// 为协程设置超时结果，并让其线程 resume 协程
						pending->cbt = CBType::Timeout;
						if (pending->co && pending->thread) {
							 pending->thread->PostTask([co = pending->co]() {
								mco_resume(co);
							});
						}
					}
					else if (pending->type == PendingType::Async) {
						auto handler_it = m_cmd_handlers.find(pending->res_cmd);
						if (handler_it != m_cmd_handlers.end()) {
							handler_it->second(pending->request, CBType::Timeout);
						}
					}
					else {
						LOG_ERR("Unknown pending type for key: {}", key.ToString());
					}
					timeout_pendings.push_back(key);
				}
			}
			for (const auto& key : timeout_pendings) {
				m_pendings.erase(key);
			}
		}
	}
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::OnConnected() 
{
	if (m_event_handler) m_event_handler(DispatcherEventType::CONNECTED);
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::OnDisconnected() 
{
	if (m_event_handler) m_event_handler(DispatcherEventType::DISCONNECTED);
}

//--------------------------------------------------------------------------------------------------
//
//--------------------------------------------------------------------------------------------------
void WsMsgDispatcher::OnMessage(const std::string& msg) 
{
	this->OnWsMessage(msg);
}

} // namespace p2pms